package com.benshou.bcss.common.mdc;

import com.benshou.bcss.common.j17common.annotation.NotNull;
import com.benshou.bcss.common.trace.TraceIdContext;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.HashMap;
import java.util.Map;

/**
 * MDC Task 装饰器
 *
 * @author <a href="mailto:liuzhenbang@benshou.com">Liu Zhenbang</a>
 */
@Slf4j
public class MDCTaskDecorator implements TaskDecorator {

    @Override
    public @NotNull Runnable decorate(@NotNull Runnable runnable) {
        var parent = MDC.getCopyOfContextMap();
        return () -> {
            var current = MDC.getCopyOfContextMap();
            var finalMap = mergeContextMap(parent, current);
            MDC.setContextMap(finalMap);
            try {
                // 初始化 TraceId
                var traceId = TraceIdContext.getOrInitialize();
                logger.debug("Current thread traceId={}", traceId);
                runnable.run();
            } finally {
                MDC.setContextMap(current);
            }
        };
    }

    /**
     * 合并上下文参数映射
     *
     * @param parent 父线程的上下文参数映射
     * @param current 当前创建的线程的上下文参数映射
     * @return 合并后的上下文参数映射
     */
    private Map<String, String> mergeContextMap(Map<String, String> parent, Map<String, String> current) {
        if (parent == null || parent.isEmpty()) {
            return current;
        }

        if (current == null || current.isEmpty()) {
            return parent;
        }

        var result = new HashMap<>(current);
        parent.forEach((k, v) ->
                result.merge(k, v, (old, late) -> {
                    throw new IllegalStateException(String.format(
                            "Duplicate key %s (attempted merging values %s and %s)",
                            k, old, late));}));
        return result;
    }
}
